import random
import logging
from typing import List

from source.RFT.utils import logger
from source.RFT.models import ScenarioMetadata

def split_scenarios(
    scenarios: List[ScenarioMetadata], 
    train_ratio: float, 
    random_seed: int
) -> List[ScenarioMetadata]:
    """
    将场景划分为训练集和测试集
    
    参数:
        scenarios: 场景元数据列表
        train_ratio: 训练集比例
        random_seed: 随机种子
        
    返回:
        更新了dataset_type字段的场景元数据列表
    """
    if not scenarios:
        logger.warning("没有场景可供划分")
        return scenarios
        
    logger.info(f"开始划分数据集: 总场景数={len(scenarios)}, 训练集比例={train_ratio}, 随机种子={random_seed}")
    
    # 设置随机种子，确保可重现性
    random.seed(random_seed)
    
    # 复制场景列表，避免修改原始数据
    scenarios_copy = scenarios.copy()
    
    # 随机打乱场景顺序
    random.shuffle(scenarios_copy)
    
    # 计算训练集的数量
    train_count = int(len(scenarios_copy) * train_ratio)
    test_count = len(scenarios_copy) - train_count
    
    # 划分训练集和测试集
    for i, scenario in enumerate(scenarios_copy):
        scenario.dataset_type = "train" if i < train_count else "test"
        
        # 使用scenario_id作为文件夹名称
        scenario.level_folder = scenario.scenario_id
    
    logger.info(f"数据集划分完成: 训练集={train_count}个场景, 测试集={test_count}个场景")
    
    return scenarios_copy